package utils

import (
    "context"
    "time"
    clientv3 "go.etcd.io/etcd/client/v3"
)

type EtcdClient struct {
    Client *clientv3.Client
}

func NewEtcdClient(endpoints []string, username, password string) (*EtcdClient, error) {
    cli, err := clientv3.New(clientv3.Config{
        Endpoints:   endpoints,
        Username:    username,
        Password:    password,
        DialTimeout: 5 * time.Second,
    })
    if err != nil {
        return nil, err
    }
    return &EtcdClient{Client: cli}, nil
}

func (c *EtcdClient) Put(key, value string) error {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    _, err := c.Client.Put(ctx, key, value)
    return err
}

func (c *EtcdClient) PutWithTTL(key, value string, ttl int64) error {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    
    // Create lease
    lease, err := c.Client.Grant(ctx, ttl)
    if err != nil {
        return err
    }
    
    _, err = c.Client.Put(ctx, key, value, clientv3.WithLease(lease.ID))
    return err
}

func (c *EtcdClient) Get(key string) (string, error) {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    resp, err := c.Client.Get(ctx, key)
    if err != nil {
        return "", err
    }
    if len(resp.Kvs) == 0 {
        return "", nil
    }
    return string(resp.Kvs[0].Value), nil
}

func (c *EtcdClient) GetWithPrefix(prefix string) ([]*clientv3.Event, error) {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    resp, err := c.Client.Get(ctx, prefix, clientv3.WithPrefix())
    if err != nil {
        return nil, err
    }
    
    var events []*clientv3.Event
    for _, kv := range resp.Kvs {
        events = append(events, &clientv3.Event{
            Type: clientv3.EventTypePut,
            Kv:   kv,
        })
    }
    return events, nil
}

func (c *EtcdClient) Delete(key string) error {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    _, err := c.Client.Delete(ctx, key)
    return err
}

func (c *EtcdClient) DeleteWithPrefix(prefix string) error {
    ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    defer cancel()
    _, err := c.Client.Delete(ctx, prefix, clientv3.WithPrefix())
    return err
}

// Close 关闭etcd客户端连接
func (c *EtcdClient) Close() error {
    if c != nil && c.Client != nil {
        return c.Client.Close()
    }
    return nil
}